//
// Created by 郝羽 on 2018/2/26.
//

#include "mixer2.hpp"

#include <map>
#include <cmath>
#include <opus.h>
#include <string>

#include "audio_recveq.h"
#include "xm_log.h"
#include "xm_app_config.h"

#define mix_dbgi dbgi
#define mix_dbgw dbgw
#define mix_dbge dbge

namespace mixer2 {

constexpr int MTU = 1500;
constexpr int CHANNEL_NB = 2;
constexpr int SAMPLE_RATE = 16000;
constexpr int PTIME = 20;
constexpr int MIX_NB_MAX = 2;
constexpr int8_t MIX_LEVEL_MIN = -60;
constexpr int OPUS_APP = OPUS_APPLICATION_AUDIO;
constexpr int8_t RTP_PT_DUMMY = 111;
constexpr uint32_t RTP_SSRC_DUMMY = 1234;
constexpr int RTP_SAMPLE_RATE = 48000;

typedef int16_t sample_t;

// as per https://tools.ietf.org/html/rfc6465#appendix-A.1
static
int8_t calc_level(sample_t* samples, size_t size) {
    if (size == 0) {
        return 0;
    }

    double rms = 0;
    for(int i = 0; i < size; i++){
        double v = static_cast<double>(samples[i]);
        v /= 32767;
        rms += v * v;
    }
    rms = std::sqrt(rms / size);
    int db = 20 * std::log10(rms);
    if (db < -127) {
        db = -127;
    } else if (db > 0) {
        db = 0;
    }
    return db;
}

static
void ramp_samples(sample_t* samples, float from, float to) {
    int nb_sample_per_ch = SAMPLE_RATE / 1000 * PTIME;
    if (nb_sample_per_ch <= 1) {
        return;
    }
    float step = (to - from) / (nb_sample_per_ch - 1);
    float gain = from;
    for(int i = 0; i < nb_sample_per_ch; i++) {
        for(int ch = 0; ch < CHANNEL_NB; ch++){
            int pos = i * CHANNEL_NB + ch;
            samples[pos] *= gain;
        }
        gain += step;
    }
}

static
int make_rtp_header(rtc_endpoint_t rtc, uint8_t* buf, uint16_t seq, uint32_t ts, std::list<stream*> streams) {
    //     0                   1                   2                   3
    //     0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
    //     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //     |V=2|P|X|  CC   |M|     PT      |       sequence number         |
    //     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //     |                           timestamp                           |
    //     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //     |           synchronization source (SSRC) identifier            |
    //     +=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
    //     |            contributing source (CSRC) identifiers             |
    //     |                             ....                              |
    //     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    uint8_t cc = 0;
    for (auto it : streams) {
        cc += 1;
        // rtp could at most have 15 csrc
        if (cc == 15) {
            break;
        }
    }

    int x = 0;
    int csrcAudioLevelId = 0;
    int ssrcAudioLevelId = 0;
    if (rtc) {
        auto ext = rtc_endpoint_audio_extension_map(rtc);
        csrcAudioLevelId = ext->getId(RtpExtensionType::CSRCAudioLevel);
        ssrcAudioLevelId = ext->getId(RtpExtensionType::SSRCAudioLevel);
        if ((csrcAudioLevelId && cc) || ssrcAudioLevelId) {
            x = 1;
        }
    }
    int offset = 0;
    buf[offset++] = 2 << 6 | 0 << 5 | x << 4 | cc;
    buf[offset++] = 1 << 7 | RTP_PT_DUMMY;
    be_set_u16(seq, buf + offset);
    offset += 2;
    be_set_u32(ts, buf + offset);
    offset += 4;
    be_set_u32(RTP_SSRC_DUMMY, buf + offset);
    offset += 4;

    // audio level ranges from -127(zero energy) to 0 (full of energy)
    int8_t max_audio_level = -127;
    int i = 0;
    for (auto it : streams) {
        if (it->level() > max_audio_level) {
            max_audio_level = it->level();
        }
        be_set_u32(it->ssrc(), buf+12+4*i);
        offset += 4;
        ++i;
        if (i == 15) {
            break;
        }
    }

    // write audio-level only if rtp extension header is present,
    // prefer csrc-audio-level against ssrc-audio-level
    //
    //   0                   1                   2                   3
    //   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
    //   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //   |       0xBE    |    0xDE       |           length=3            |
    //   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //   |  ID   | L=0   |     data      |  ID   |  L=1  |   data...
    //   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //         ...data   |    0 (pad)    |    0 (pad)    |  ID   | L=3   |
    //   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    //   |                          data                                 |
    //   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    if (csrcAudioLevelId && cc) {
        buf[offset++] = 0xBE;
        buf[offset++] = 0xDE;
//            auto n = csrcs.size() - (exclude_csrc ? 1 : 0);
        auto extLen = 1 + cc;
        extLen = (extLen + 3) / 4;
        be_set_u16(extLen, buf + offset);
        offset += 2;
        //   0                   1                   2                   3
        //   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
        //   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        //   |  ID   | len=2 |0|   level 1   |0|   level 2   |0|   level 3   |
        //   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        //
        // here we used type of the extension, which get translated by send_rtp_data
        //buf[offset++] = csrcAudioLevelId << 4 | (n - 1);
        buf[offset++] = RtpExtensionType::CSRCAudioLevel << 4 | (cc - 1);
        int i = 0;
        for (auto it : streams) {
            buf[offset++] = it->level() | 0x80;
            i++;
            if (i == 15) {
                break;
            }
        }
    } else if (ssrcAudioLevelId){
        buf[offset++] = 0xBE;
        buf[offset++] = 0xDE;
        be_set_u16(1, buf + offset);
        offset += 2;
        //    0                   1
        //    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5
        //    +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        //    |  ID   | len=0 |V| level       |
        //    +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        //
        // here we used type of the extension, which get translated by send_rtp_data
        // buf[offset++] = ssrcAudioLevelId << 4 | 0;
        buf[offset++] = RtpExtensionType::SSRCAudioLevel << 4 | 0;
        buf[offset++] = -max_audio_level | 0x80;
    }
    if (csrcAudioLevelId || ssrcAudioLevelId) {
        while (offset % 4) {
            buf[offset++] = 0;
        }
    }

    return offset;
}

class stream {
private:
    std::string mixer_id_;
    rtc_endpoint_t rtc_;
    uint8_t encoded_data_[MTU] = {0,};
    sample_t decoded_data_[SAMPLE_RATE / 1000 * PTIME * CHANNEL_NB] = {0,};
    size_t decoded_data_len_ = 0;
    OpusEncoder* encoder_ = nullptr;
    int8_t level_ = -127;
    audio_recveq_t eq_ = nullptr;
    // set true if selected for this round, clear after mixed
    bool mix_ = false;
    // set to 1.0 after mixed, set to 0.0 after not mixed
    float gain_ = 0.0;

public:
    stream(const std::string& mixer_id, rtc_endpoint_t rtc)
            :mixer_id_(mixer_id), rtc_(rtc) {}

    int init() {
        int err = 0;
        encoder_ = opus_encoder_create(SAMPLE_RATE, CHANNEL_NB, OPUS_APP, &err);
        if (err) {
            return err;
        }
        opus_encoder_ctl(encoder_, OPUS_SET_COMPLEXITY_REQUEST, 0);
        opus_encoder_ctl(encoder_, OPUS_SET_INBAND_FEC_REQUEST, 1);
        opus_encoder_ctl(encoder_, OPUS_SET_PACKET_LOSS_PERC_REQUEST, 20);
        int opus_lsb = app_config_get()->mixer_opus_lsb;
        if (opus_lsb >= 8 && opus_lsb < 16) {
            opus_encoder_ctl(encoder_, OPUS_SET_LSB_DEPTH(opus_lsb));
            opus_encoder_ctl(encoder_, OPUS_SET_DTX(1));
        }
        eq_ = audio_recveq_create();

        return 0;
    }

    bool push(unsigned char* buf, size_t size, int8_t min_level) {
        int vad = 0;
        int level = 0;
        audio_recveq_prepush(eq_, buf, size, &level, &vad);
        if (level >= 0) {
            level_ = -level;
        } else {
            level_ = -127;
        }
        // todo: also commit if mixed in last round, so that we can fade out
        if (level_ > min_level) {
            audio_recveq_commit_push(eq_);
            return true;
        }
        return false;
    }

    int8_t level() {
        return level_;
    }

    void flush() {
        audio_recveq_flush(eq_);
    }

    int pull(sample_t* out) {
        uint32_t ts;
        sample_t* sample;
        int frame_size;
        int channel_nb;

        decoded_data_len_ = 0;
        for (int t = 0; t < PTIME; t += 10) {
            if (int err = audio_recveq_pull_pcm(eq_, &ts, &sample, &frame_size, &channel_nb)) {
                dbge("mixer |%s| stream |%s| pull failed", mixer_id_.c_str(), rtc_endpoint_get_id(rtc_));
                return err;
            }
            int pulled = frame_size * channel_nb * sizeof(sample_t);
            memcpy(decoded_data_ + decoded_data_len_, sample, pulled);
            decoded_data_len_ += pulled;
        }

        mix_in(out);
        return 0;
    }

    void send(sample_t* samples, uint8_t* encoded, size_t encoded_size, std::vector<stream*> streams, uint16_t rtp_seq, uint32_t rtp_ts) {
        int rtp_hdr_len = make_rtp_header(rtc_, encoded_data_, rtp_seq, rtp_ts, streams);
        // substract own contribution from samples if mixed
        if (state_ == mix_state::FADING_IN || state_ == mix_state::ACTIVE) {
            undo_mix_in(samples);
            int frame_size = SAMPLE_RATE / 1000 * PTIME;
            int rtp_bdy_len = opus_encode(encoder_, decoded_data_, frame_size, encoded_data_ + rtp_hdr_len, sizeof(encoded_data_) - rtp_hdr_len);
            rtc_endpoint_send_stream_data(rtc_, RTC_AUDIO, rtc_codec_id_opus, rtc_codec_id_unknown, 0, 0, encoded_data_, rtp_hdr_len + rtp_bdy_len, nullptr);
        } else {
            memcpy(encoded_data_ + rtp_hdr_len, encoded, encoded_size);
            rtc_endpoint_send_stream_data(rtc_, RTC_AUDIO, rtc_codec_id_opus, rtc_codec_id_unknown, 0, 0, encoded_data_, rtp_hdr_len + encoded_size, nullptr);
        }
    }

    void warmup(sample_t* samples) {
        if (state_ == mix_state::FADING_IN) {
            int frame_size = SAMPLE_RATE / 1000 * PTIME;
            opus_encode(encoder_, samples, frame_size, encoded_data_, sizeof(encoded_data_));
        }
    }

private:
    // todo: need some protection from saturation
    void mix_in(sample_t* out) {
        sample_t* in = decoded_data_;
        int n = SAMPLE_RATE * PTIME;
        for (int i = 0; i < n; ++i) {
            for (int j = 0; j < CHANNEL_NB; ++j) {
                int k = i * CHANNEL_NB + j;
                out[k] += in[k];
            }
        }
    }

    // todo: need some protection from saturation
    void undo_mix_in(sample_t* in) {
        sample_t* out = decoded_data_;
        int n = SAMPLE_RATE * PTIME;
        for (int i = 0; i < n; ++i) {
            for (int j = 0; j < CHANNEL_NB; ++j) {
                int k = i * CHANNEL_NB + j;
                // out will be inversed, no impact on audio perception
                out[k] -= in[k];
            }
        }
    }
};

class mixer {
private:
    std::map<rtc_endpoint_t, stream*> streams_;
    // todo: fix this !!
    // candidate keeps tracks of stream mixed last round,
    // if it has a higher level, and stops to update,
    // new candidate will fail to get in,
    // so either we invalidate old candidate faster
    // or we can't save old candidate here
    std::vector<stream*> candidates_;
    std::vector<stream*> prev_candidates_;
    std::string id_;
    OpusEncoder* encoder_;
    uint8_t encoded_data_[MTU];
    sample_t decoded_data_[SAMPLE_RATE / 1000 * PTIME * CHANNEL_NB];
    uint32_t rtp_ts_ = 0;
    uint16_t rtp_seq_ = 10000;
public:
    mixer(const std::string& id)
            :id_(id)
    {
    }

    int add(rtc_endpoint_t rtc) {
        auto s = new stream(id_, rtc);
        if (int err = s->init()) {
            return err;
        }
        streams_.emplace(rtc, s);
    }

    int push(rtc_endpoint_t rtc, unsigned char* buf, size_t size) {
        auto it = streams_.find(rtc);
        if (it == streams_.end()) {
            return -2;
        }
        auto s = it->second;
        int8_t min_level = MIX_LEVEL_MIN;

        if (candidates_.size() >= MIX_NB_MAX) {
            std::sort(candidates_.begin(), candidates_.end(),
                      [](const stream* a, const stream* b) {
                return a < b;
            });
            auto min = *candidates_.begin();
            min_level = (*candidates_.begin())->level();
        }

        if (s->push(buf, size, min_level)) {
            if (candidates_.size() >= MIX_NB_MAX) {
                auto old = *candidates_.begin();
                old->flush();
                s->flush();
                *candidates_.begin() = s;
            } else {
                candidates_.emplace_back(s);
            }
        }
        return 0;
    }

    int mix() {
        clear_mix_buf();
        for (auto s : candidates_) {
            s->warmup(decoded_data_);
        }
        for (auto s : candidates_) {
            s->pull(decoded_data_);
        }
        int frame_size = SAMPLE_RATE / 1000 * PTIME;
        int n = opus_encode(encoder_, decoded_data_, frame_size, encoded_data_, MTU);
        if (n < 0) {
            return -1;
        }
        for (auto it : streams_) {
            auto s = it.second;
            s->send(decoded_data_, encoded_data_, n, candidates_, rtp_seq_, rtp_ts_);
        }
        rtp_seq_ += 1;
        rtp_ts_ += RTP_SAMPLE_RATE / 1000 * PTIME;
    }
private:
    void clear_mix_buf() {
        int n = SAMPLE_RATE * PTIME;
        for (int i = 0; i < n; ++i) {
            for (int j = 0; j < CHANNEL_NB; ++j) {
                int k = i * CHANNEL_NB + j;
                decoded_data_[k] = 0;
            }
        }
    }
};

}

void* mixer2_create(const char* id) {
    return new mixer2::mixer(id);
}

void mixer2_destroy(void* p_void) {
    auto mixer = static_cast<mixer2::mixer*>(p_void);
    delete mixer;
}

int mixer2_add(void* p_void, rtc_endpoint_t rtc) {
    auto mixer = static_cast<mixer2::mixer*>(p_void);
    if (!mixer) {
        return -1;
    }
    mixer->add(rtc);
}

int mixer2_push(void* p_void, rtc_endpoint_t rtc, unsigned char* buf, size_t size) {
    auto mixer = static_cast<mixer2::mixer*>(p_void);
    if (!mixer) {
        return -1;
    }
    return mixer->push(rtc, buf, size);
}

